#include "codec_h264.h"

static int test_open_files(char **argv)
{
	const char* inputFileName = argv[1];
	const char* outputFileName = argv[2];

	pFin = fopen(inputFileName, "rb+");
	if (!pFin)
	{
		printf("Error: open input file failed.\n");
		return -1;
	}
	pFout = fopen(outputFileName, "wb+");
	if (!pFout)
	{
		printf("Error: open input file failed.\n");
		return -1;
	}
	return 0;
}

static void test_write_yuv_file(AVFrame *frame)
{
	uint8_t **pBuf = frame->data;
	int *pStride = frame->linesize;

	for (int color_idx = 0; color_idx < 3; color_idx++)
	{
		int	nWidth = color_idx == 0 ? frame->width : frame->width / 2;
		int	nHeight = color_idx == 0 ? frame->height : frame->height / 2;
		for (int idx = 0; idx < nHeight; idx++)
		{
			fwrite(pBuf[color_idx], 1, nWidth, pFout);
			pBuf[color_idx] += pStride[color_idx];
		}
		fflush(pFout);
	}
}

static void test_close_files(void)
{
	fclose(pFin);
	fclose(pFout);
}

void test_wirte_h264_file(AVPacket *pkt)
{
	printf("The pkt size is %d\n", pkt->size);
	fwrite(pkt->data, 1, pkt->size, pFout);
}

void test_codec_h264_decoder_app_main(const char **argv)
{
	uint8_t inbuf[INBUF_SIZE + AV_INPUT_BUFFER_PADDING_SIZE];

	test_open_files(argv);

	struct ffmpeg_player player = {};
	codec_h264_init_ffmpeg_player(&player, AV_CODEC_ID_H264);

	codec_h264_config_decoder_player(&player);
	int uDataSize = 0;

	while (1)
	{
		uDataSize = fread(inbuf, 1, INBUF_SIZE, pFin);
		if (uDataSize == 0)
		{
			break;
		}
		if (uDataSize > 0)
			codec_h264_decode_one_frame(&player, inbuf, uDataSize, 20);
		else
			break;
	}
	codec_h264_deinit_ffmpeg_player(&player);
	test_close_files();
}

void test_codec_h264_encoder_app_main(const char **argv)
{
	struct ffmpeg_player player = {};
#define WIDHT 600
#define HEIGHT 316

	test_open_files(argv);

	codec_h264_init_ffmpeg_player(&player, AV_CODEC_ID_H264);

	codec_h264_config_encodec_player(&player, WIDHT, HEIGHT, 30, AV_PIX_FMT_NV12);

	int size = WIDHT * HEIGHT * 3 / 2;

	while (1) {
		int len = fread(player.frame_fmt, 1, size, pFin);
		if (len <= 0) {
			printf("touch the end of file\n");
		break;
	}

	printf("encode %d\n", len);
	/* converted data */
	player.what_frame->data[0] = player.frame_fmt;
	player.what_frame->data[1] = player.frame_fmt + WIDHT * HEIGHT;
	player.what_frame->data[2] = player.frame_fmt + WIDHT * HEIGHT * 5 /4;

	int ret = codec_h264_encode_one_frame(&player);
	if (-1 == ret) {
		printf("encode one frame error\n");
		break;
		}
	}
	test_close_files();
}

int codec_h264_init_ffmpeg_player(struct ffmpeg_player *player, int codec_id)
{
	avcodec_register_all();
	av_init_packet(&player->pkt);

	player->codec = avcodec_find_decoder(codec_id);
	if (!player->codec)
	{
		printf("find decoder error\n");
		return -1;
	}

	player->codec_ctx = avcodec_alloc_context3(player->codec);
	if (!player->codec_ctx)
	{
		printf("alloc context error\n");
		return -1;
	}

	player->coder_parser_ctx = av_parser_init(codec_id);
	if (!player->coder_parser_ctx)
	{
		printf("alloc parser error.\n");
		return -1;
	}

	return 0;
}

int codec_h264_config_encodec_player(struct ffmpeg_player *player, int width, int height, int fps, int raw_format)
{
	player->codec_ctx->bit_rate = 400000;
	player->codec_ctx->width = width;
	player->codec_ctx->height = height;
	player->codec_ctx->time_base = (AVRational) {1, fps};
	player->codec_ctx->framerate = (AVRational) {fps, 1};
	player->codec_ctx->gop_size = 10;
	player->codec_ctx->max_b_frames = 1;
	player->codec_ctx->pix_fmt = raw_format;

	if (avcodec_open2(player->codec_ctx, player->codec, NULL) < 0)
	{
		printf("Error:open condec failed.\n");
		return -1;
	}


	player->what_frame = av_frame_alloc();
	if (!player->what_frame)
	{
		printf("Error:alloc frame failed.\n");
		return -1;
	}

	player->frame_fmt = (unsigned char *)av_malloc(av_image_get_buffer_size(raw_format, width, height, 1));
	if (!player->frame_fmt) {
		printf("av alloc error");
		return -1;
	}

	av_image_fill_arrays(player->what_frame->data, player->what_frame->linesize, raw_format, width, height, 1);

	player->what_frame->width = width;
	player->what_frame->height = height;
	player->what_frame->format = raw_format;

	return 0;
}

int codec_h264_config_decoder_player(struct ffmpeg_player *player)
{
	if (avcodec_open2(player->codec_ctx, player->codec, NULL) < 0)
	{
		printf("Error:open condec failed.\n");
		return -1;
	}

	player->what_frame = av_frame_alloc();
	if (!player->what_frame)
	{
		printf("Error:alloc frame failed.\n");
		return -1;
	}
	return 0;
}

void codec_h264_deinit_ffmpeg_player(struct ffmpeg_player *player)
{
	avcodec_close(player->codec_ctx);
	av_free(player->codec_ctx);
	av_frame_free(&player->what_frame);
}

int codec_h264_decode_one_frame(struct ffmpeg_player *player, unsigned char *raw_data, int raw_data_len, int timeout_cnt)
{
	int err;
	int parser_len;
	unsigned char *raw_data_ptr = NULL;
	int got_frame;
	raw_data_ptr = raw_data;
	static int index_frame;

	while (raw_data_len > 0)
	{
		parser_len = av_parser_parse2(player->coder_parser_ctx, player->codec_ctx,
		&player->pkt.data, &player->pkt.size,
		raw_data_ptr, raw_data_len, AV_NOPTS_VALUE, AV_NOPTS_VALUE, AV_NOPTS_VALUE);
		raw_data_ptr += parser_len;
		raw_data_len -= parser_len;

		if (player->pkt.size == 0) {
			if (!timeout_cnt--)
				break;
			continue;
		}

		err = avcodec_decode_video2(player->codec_ctx, player->what_frame, &got_frame, &player->pkt);
		if (err < 0)
		{
			printf("Error: decode failed.\n");
			return -1;
		}

		if (!got_frame)
			break;

		printf("decoded one frame index frame %d, width x height: (%d x %d)\n", index_frame++, player->what_frame->width, player->what_frame->height);
		test_write_yuv_file(player->what_frame);
	}

	player->pkt.data = NULL;
	player->pkt.size = 0;

	return 0;
}

int codec_h264_encode_one_frame(struct ffmpeg_player *player)
{
#if 0 //deprecated API
	int err;
	int got_picture;

	err = avcodec_encode_video2(player->codec_ctx, &player->pkt, player->what_frame, &got_picture);
	if (err < 0) {
		printf("encode error");
		return -1;
	}

	if (got_picture) {
		test_wirte_h264_file(&player->pkt);
		printf("got picture to packet\n");
	}
	av_free_packet(&player->pkt);
	#endif
	int ret = -1;
	
	/* from frame to packet for encoder */
	avcodec_send_frame(player->codec_ctx, player->what_frame);

	while (avcodec_receive_packet(player->codec_ctx, &player->pkt) >= 0) {
		test_wirte_h264_file(&player->pkt);
	}
	return 0;
}